"""
model name : aihw
file       : astar.py
information:
    author : OuYang
    time   : 2025/10/10
"""
import copy
import heapq
from utils.node import Node
from utils.utils import corr2index


class Astar:
    def __init__(self, basemap, col=19, row=19):
        self.col = col
        self.row = row
        self.basemap = basemap

        self.start_node = Node(1, 1, 0)
        self.end_node = Node(self.col, self.row, 0)

    @staticmethod
    def heuristic(start_node, end_node):
        return abs(end_node.x - start_node.x) + abs(end_node.y - start_node.y)

    @staticmethod
    def reconstruct_path(current_node: Node):
        path = []
        all_weights = 0
        current = current_node

        while current is not None:
            all_weights += current.w
            path.append((current.x, current.y))
            current = current.parent

        path.reverse()
        return path, all_weights

    def set_start_node(self, x, y):
        self.start_node = Node(x, y, 0)

    def set_end_node(self, x, y):
        self.end_node = Node(x, y, 0)

    def get_neighbors(self, node):

        neighbors = []
        directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]

        for dx, dy in directions:
            nx, ny = node.x + dx, node.y + dy
            if 0 < nx <= self.col and 0 < ny <= self.row:
                i, j = corr2index(nx, ny)
                if self.basemap[i][j] != -1:
                    neighbors.append(Node(nx, ny, self.basemap[i][j]))

        return neighbors

    def solve(self):

        # 初始化起点
        start_node = copy.deepcopy(self.start_node)
        start_node.g = float(0)
        start_node.h = self.heuristic(self.start_node, self.end_node)
        start_node.f = start_node.g + start_node.h

        # 开放列表和关闭列表
        open_list = []
        close_set = set()

        heapq.heappush(open_list, start_node)

        # 用于记录所有节点的字典，便于查找和更新
        all_nodes = {(start_node.x, start_node.y): start_node}

        # 用于记录搜索过程
        search_steps = []

        while open_list:
            # 获取当前优先队列中f值最小的元素
            current_node = heapq.heappop(open_list)

            search_steps.append((current_node.x, current_node.y))

            # 检测是否到达终点
            if current_node == self.end_node:
                path, all_weights = self.reconstruct_path(current_node)
                return path, search_steps, all_weights

            # 未达终点，将当前节点加入关闭队列
            close_set.add((current_node.x, current_node.y))

            # 检查所有邻居节点
            for neighbor in self.get_neighbors(current_node):
                neighbor_corr = neighbor.x, neighbor.y

                if neighbor_corr in close_set:
                    continue

                # 计算从起点经过当前节点到达邻居节点
                tentative_g = current_node.g + neighbor.w

                # 如何这个邻居节点不在以访问过的全部节点，
                # 或者通过当前节点到达邻居节点的总花费小于通过其他路径到达邻居节点的花费(找到了更优路径)
                if neighbor_corr not in all_nodes or tentative_g < all_nodes[neighbor_corr].g:
                    if neighbor_corr not in all_nodes:
                        neighbor_node = copy.deepcopy(neighbor)
                        all_nodes[neighbor_corr] = neighbor_node
                    else:
                        neighbor_node = all_nodes[neighbor_corr]

                    neighbor_node.parent = current_node
                    neighbor_node.g = tentative_g
                    neighbor_node.h = self.heuristic(neighbor_node, self.end_node)
                    neighbor_node.f = neighbor_node.g + neighbor_node.h

                    # 如果邻居节点不在开放列表中
                    if neighbor_node not in open_list:
                        heapq.heappush(open_list, neighbor_node)
                    else:
                        heapq.heapify(open_list)

        return None, search_steps, -1
